package com.jorado.pulsar.client.trace;

import brave.internal.HexCodec;
import brave.propagation.TraceContext;
import com.jorado.json.JsonUtils;
import com.jorado.trace.util.RequestID;
import org.apache.pulsar.client.api.Message;
import org.apache.pulsar.client.api.MessageId;
import org.apache.pulsar.client.api.Producer;
import org.apache.pulsar.client.api.ProducerInterceptor;
import org.apache.pulsar.client.impl.MessageImpl;
import org.apache.pulsar.common.api.proto.PulsarApi;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.HashMap;
import java.util.Map;

public class TracingProducerInterceptor<T> implements ProducerInterceptor<T> {

    private static final Logger LOG = LoggerFactory.getLogger(TracingProducerInterceptor.class);
    private String clientAppName = "pulsar";

    private TraceContextContainer traceContextContainer;

    public TracingProducerInterceptor(String clientAppName) {
        traceContextContainer = new TraceContextContainer("onSendAcknowledgement", clientAppName);
        this.clientAppName = clientAppName;

    }

    @Override
    public void close() {

    }

    @Override
    public Message<T> beforeSend(Producer<T> producer, Message<T> message) {
        if (message != null) {
            messageSetProperties(message);
        }
        return message;
    }

    @Override
    public void onSendAcknowledgement(Producer<T> producer, Message<T> message, MessageId messageId, Throwable throwable) {
        writeResultLog(producer.getTopic(), messageId, throwable, message);
    }

    private void messageSetProperties(Message<T> message) {
        TraceContext traceContext = traceContextContainer.initZipkinContextWithoutScope();
        if (message instanceof MessageImpl) {
            PulsarApi.MessageMetadata.Builder builder = ((MessageImpl) message).getMessageBuilder();
            builder.addProperties(PulsarApi.KeyValue.newBuilder().setValue(RequestID.getRequestID()).setKey("requestId"));
            builder.addProperties(PulsarApi.KeyValue.newBuilder().setValue(traceContext.spanId() + "").setKey("spanId"));
            builder.addProperties(PulsarApi.KeyValue.newBuilder().setValue(traceContext.parentId() + "").setKey("parentId"));
            builder.addProperties(PulsarApi.KeyValue.newBuilder().setValue(System.currentTimeMillis() + "").setKey("startTime"));
            builder.addProperties(PulsarApi.KeyValue.newBuilder().setValue(traceContext.sampled() + "").setKey("sampled"));
        } else {
            LOG.error("Can't set properties on TopicMessage.");
        }
    }

    protected void writeResultLog(String topicName, MessageId messageId, Throwable throwable, Message<T> message) {
        Map invLog = new HashMap<>();
        invLog.put("type", "ss");
        invLog.put("service", "PulsarProducerClient");
        invLog.put("method", "send");
        invLog.put("spanId", HexCodec.toLowerHex(Long.parseLong(message.getProperty("spanId"))));
        invLog.put("parentId", HexCodec.toLowerHex(Long.parseLong(message.getProperty("parentId"))));
        invLog.put("topicName", topicName);
        RequestID.setRequestID(message.getProperty("requestId"));
        invLog.put("traceId", RequestID.getRequestID());
        invLog.put("messageId", messageId.toString());
        invLog.put("pulsarTags", "pulsar");
        invLog.put("serverType", "httpService");
        invLog.put("Params", "");
        invLog.put("ExecuteTime", System.currentTimeMillis() - Long.parseLong(message.getProperty("startTime")));
        invLog.put("project", clientAppName);
        if (throwable == null) {
            invLog.put("errorCode", "success");
        } else {
            invLog.put("errorCode", throwable);
        }
        LOG.info(JsonUtils.toJson(invLog));
    }

    public TraceContextContainer getTraceContextContainer() {
        return traceContextContainer;
    }
}

